#!/usr/bin/env python

import os, sys
import numpy as np
from astropy.io import fits
from astropy.visualization import ZScaleInterval
import matplotlib.pylab as plt
from scipy.stats import sigmaclip


def get_prnu(image, varmap):
    var_tot = sigmaclip(image, 5, 5)[0].var(ddof=1)
    var_pix = np.median(varmap)*0
    print('vat_rot = {}'.format(var_tot))
    print('var_pix = {}'.format(var_pix))
    return np.sqrt(var_tot-var_pix)/image.mean()

def get_prnu_map(image, varmap):
    return image / varmap

zs = ZScaleInterval()

flat_e = fits.open(sys.argv[1])[0].data
noise_e = fits.open(sys.argv[1])[1].data

model_e = fits.open(sys.argv[2])[0].data

nx, ny = flat_e.shape[0], flat_e.shape[1]

###################################################
# replace OS4 with OS13
#flat_e_os4 = flat_e[(nx//2):, 1152*3:1152*4]
#model_e_os4 = model_e[(nx//2):, 1152*3:1152*4]
#flat_e_os4 = np.flip(flat_e_os4,axis=0)
#model_e_os4 = np.flip(model_e_os4, axis=0)
#
#flat_e[0:nx//2, 1152*3:1152*4] = flat_e_os4
#model_e[0:nx//2, 1152*3:1152*4] = model_e_os4
###################################################

# calculate PRNU

prnu = get_prnu(flat_e, noise_e**2) * 100

mod = model_e / np.mean(model_e)
prnu_corr = get_prnu(flat_e/mod, noise_e**2/(mod**2))*100

print('PRNU: {}'.format(prnu))
print('PRNU_corr: {}'.format(prnu_corr))


#data = flat_e[0:nx//2, 0:ny//2] / model_e[0:nx//2, 0:ny//2]
data = flat_e / model_e
index = np.isfinite(data)
data = data / data.mean()

fig = plt.figure(figsize=(6,9))

# plot PRNU map
ax = fig.add_subplot(2,1,1)
vmin, vmax = zs.get_limits(data[index])
vmin, vmax = 1.05, 0.975
plt.imshow(data, vmin=vmin, vmax=vmax, cmap='Blues_r', origin='lower')
plt.colorbar()
#plt.title('300nm')
#plt.savefig('prnu.png')

ax = fig.add_subplot(2,1,2)
#plt.figure()
prnu_sigma = np.std(sigmaclip(data, 5, 5)[0])
print('prnu_sigma: {}'.format(prnu_sigma))
prnu_sigma_lb = 'sigma: {:.3f}'.format(prnu_sigma)
plt.hist(data.flatten(), bins=100, range=[0.9,1.1], density=True,
         label=prnu_sigma_lb)
plt.xlabel('Dist of normed pixel values')
plt.legend(loc='upper left')
plt.tight_layout()
plt.show()
